#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI营销方案生成系统 - WebSocket模块
提供实时通信功能和事件处理

主要功能：
- WebSocket连接管理
- 实时消息推送
- 聊天对话支持
- 生成进度通知
- 系统状态广播

设计原则：
- 事件驱动：基于事件的消息处理
- 房间管理：支持用户分组和广播
- 状态同步：保持客户端状态一致
- 错误恢复：处理连接断开和重连

Author: AI Marketing System
Date: 2024
"""

from flask_socketio import SocketIO, emit, join_room, leave_room, disconnect
from flask import request, session
import logging

# 导入事件处理器
from .events.connection import ConnectionHandler
from .events.chat import ChatHandler
from .events.generation import GenerationHandler
from .events.notification import NotificationHandler
from .events.system import SystemHandler

# 导入工具类
from .utils.room_manager import RoomManager
from .utils.session_manager import SessionManager
from .utils.message_queue import MessageQueue
from .utils.rate_limiter import WebSocketRateLimiter

# 全局SocketIO实例
socketio = None

# 事件处理器注册表
EVENT_HANDLERS = {
    'connection': ConnectionHandler,
    'chat': ChatHandler,
    'generation': GenerationHandler,
    'notification': NotificationHandler,
    'system': SystemHandler
}

# 全局管理器实例
room_manager = RoomManager()
session_manager = SessionManager()
message_queue = MessageQueue()
rate_limiter = WebSocketRateLimiter()

# 日志记录器
logger = logging.getLogger(__name__)


class WebSocketManager:
    """
    WebSocket管理器
    负责WebSocket连接和事件的统一管理
    """
    
    def __init__(self, app=None, socketio_instance=None):
        """
        初始化WebSocket管理器
        
        Args:
            app: Flask应用实例
            socketio_instance: SocketIO实例
        """
        self.app = app
        self.socketio = socketio_instance
        self.handlers = {}
        self.middleware = []
        
        if app and socketio_instance:
            self.init_app(app, socketio_instance)
    
    def init_app(self, app, socketio_instance):
        """
        初始化Flask应用和SocketIO
        
        Args:
            app: Flask应用实例
            socketio_instance: SocketIO实例
        """
        self.app = app
        self.socketio = socketio_instance
        
        # 初始化管理器
        room_manager.init_app(app)
        session_manager.init_app(app)
        message_queue.init_app(app)
        rate_limiter.init_app(app)
        
        # 注册事件处理器
        self.register_handlers()
        
        # 注册中间件
        self.register_middleware()
        
        logger.info("WebSocket管理器初始化完成")
    
    def register_handlers(self):
        """
        注册所有事件处理器
        """
        for handler_name, handler_class in EVENT_HANDLERS.items():
            handler_instance = handler_class(
                self.socketio,
                room_manager,
                session_manager,
                message_queue
            )
            
            self.handlers[handler_name] = handler_instance
            handler_instance.register_events()
            
            logger.info(f"事件处理器 '{handler_name}' 注册成功")
    
    def register_middleware(self):
        """
        注册WebSocket中间件
        """
        # 认证中间件
        @self.socketio.on('connect')
        def handle_connect(auth=None):
            """
            处理客户端连接
            """
            try:
                # 限流检查
                client_ip = request.environ.get('REMOTE_ADDR')
                if not rate_limiter.allow_connection(client_ip):
                    logger.warning(f"连接被限流拒绝: {client_ip}")
                    return False
                
                # 认证检查
                if not self._authenticate_connection(auth):
                    logger.warning(f"连接认证失败: {client_ip}")
                    return False
                
                # 记录连接
                session_manager.add_session(request.sid, {
                    'user_id': session.get('user_id'),
                    'ip': client_ip,
                    'user_agent': request.headers.get('User-Agent'),
                    'connected_at': session_manager.get_current_timestamp()
                })
                
                logger.info(f"客户端连接成功: {request.sid}")
                return True
                
            except Exception as e:
                logger.error(f"连接处理异常: {e}")
                return False
        
        # 断开连接中间件
        @self.socketio.on('disconnect')
        def handle_disconnect():
            """
            处理客户端断开连接
            """
            try:
                # 清理会话
                session_info = session_manager.get_session(request.sid)
                if session_info:
                    # 离开所有房间
                    user_rooms = room_manager.get_user_rooms(session_info.get('user_id'))
                    for room in user_rooms:
                        room_manager.leave_room(session_info.get('user_id'), room)
                    
                    # 移除会话
                    session_manager.remove_session(request.sid)
                
                logger.info(f"客户端断开连接: {request.sid}")
                
            except Exception as e:
                logger.error(f"断开连接处理异常: {e}")
        
        # 错误处理中间件
        @self.socketio.on_error_default
        def default_error_handler(e):
            """
            默认错误处理器
            """
            logger.error(f"WebSocket错误: {e}")
            emit('error', {
                'message': '服务器内部错误',
                'code': 'INTERNAL_ERROR'
            })
    
    def _authenticate_connection(self, auth):
        """
        认证WebSocket连接
        
        Args:
            auth: 认证信息
            
        Returns:
            bool: 认证是否成功
        """
        # 这里可以实现具体的认证逻辑
        # 例如验证JWT token、session等
        return True
    
    def broadcast_to_room(self, room, event, data):
        """
        向房间广播消息
        
        Args:
            room: 房间名称
            event: 事件名称
            data: 消息数据
        """
        if self.socketio:
            self.socketio.emit(event, data, room=room)
    
    def send_to_user(self, user_id, event, data):
        """
        向特定用户发送消息
        
        Args:
            user_id: 用户ID
            event: 事件名称
            data: 消息数据
        """
        user_sessions = session_manager.get_user_sessions(user_id)
        for session_id in user_sessions:
            if self.socketio:
                self.socketio.emit(event, data, room=session_id)
    
    def get_stats(self):
        """
        获取WebSocket统计信息
        
        Returns:
            dict: 统计信息
        """
        return {
            'active_connections': session_manager.get_active_count(),
            'active_rooms': room_manager.get_room_count(),
            'message_queue_size': message_queue.get_queue_size(),
            'handlers': list(self.handlers.keys()),
            'rate_limit_stats': rate_limiter.get_stats()
        }


# 全局WebSocket管理器实例
websocket_manager = WebSocketManager()


def init_websocket(app, socketio_instance):
    """
    初始化WebSocket系统
    
    Args:
        app: Flask应用实例
        socketio_instance: SocketIO实例
        
    Returns:
        WebSocketManager: WebSocket管理器实例
    """
    global socketio
    socketio = socketio_instance
    
    websocket_manager.init_app(app, socketio_instance)
    return websocket_manager


def get_websocket_manager():
    """
    获取WebSocket管理器实例
    
    Returns:
        WebSocketManager: WebSocket管理器实例
    """
    return websocket_manager


def get_socketio():
    """
    获取SocketIO实例
    
    Returns:
        SocketIO: SocketIO实例
    """
    return socketio


# 导出的类和函数
__all__ = [
    # 主要类
    'WebSocketManager',
    
    # 事件处理器
    'ConnectionHandler',
    'ChatHandler',
    'GenerationHandler',
    'NotificationHandler',
    'SystemHandler',
    
    # 工具类
    'RoomManager',
    'SessionManager',
    'MessageQueue',
    'WebSocketRateLimiter',
    
    # 全局实例
    'websocket_manager',
    'room_manager',
    'session_manager',
    'message_queue',
    'rate_limiter',
    'socketio',
    
    # 工具函数
    'init_websocket',
    'get_websocket_manager',
    'get_socketio',
    
    # 常量
    'EVENT_HANDLERS'
]

# 版本信息
__version__ = '1.0.0'
__author__ = 'AI Marketing System'
__description__ = 'WebSocket module for AI Marketing System'